package com.example.module.ai.api;

import cn.hutool.json.JSONArray;
import cn.hutool.json.JSONObject;
import cn.hutool.json.JSONUtil;
import com.example.framework.security.core.filter.TokenAuthenticationFilter;
import com.example.framework.websocket.core.security.LoginUserHandshakeInterceptor;
import com.example.module.ai.config.SparkAiProperties;
import com.example.module.ai.spark.dto.GptRespDto;
import com.example.module.ai.spark.util.SparkAiUtils;
import com.example.module.ai.spark.websocket.SparkWebSocketMessageListener;
import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import okhttp3.WebSocket;
import org.springframework.stereotype.Service;

import java.util.ArrayList;
import java.util.List;
import java.util.UUID;
import java.util.concurrent.locks.Condition;
import java.util.concurrent.locks.Lock;

@Slf4j
@Service
public class SparkAiWebSocketAiApiImpl implements AiApi {

    @Resource
    private SparkAiUtils sparkAiUtils;

    @Resource
    private SparkAiProperties sparkAiProperties;

    /**
     * 对话历史存储集合，原本想做成用户id-历史会话的map，但在建立了websocket连接后，后续请求都不会经过 {@link TokenAuthenticationFilter}
     * 因此从 SecurityFrameworkUtils.getLoginUserId()是获取不到用户信息的
     * 这也是 {@link LoginUserHandshakeInterceptor} 的意义所在
     */
    private final List<GptRespDto.RoleContent> historyList = new ArrayList<>();

    @Override
    public String completion(String question) {
        SparkWebSocketMessageListener sparkWebSocketMessageListener = new SparkWebSocketMessageListener(historyList);

        WebSocket webSocket = sparkAiUtils.createWebSocket(sparkWebSocketMessageListener);
        sendMsg(webSocket, question);
        // 等待回答完成
        waitAnswer(sparkWebSocketMessageListener);

        if (sparkWebSocketMessageListener.isErrFlag()) {
            return String.format("星火大模型响应异常：%s", sparkWebSocketMessageListener.getErrMsg());
        } else {
            String totalAnswer = sparkWebSocketMessageListener.getTotalAnswer().toString();
            log.info("星火返回答案：{}", totalAnswer);
            return totalAnswer;
        }
    }

    private void sendMsg(WebSocket webSocket, String question) {
        JSONObject requestJson = new JSONObject();

        JSONObject header = new JSONObject();  // header参数
        header.set("app_id", sparkAiProperties.getAppId());
        header.set("uid", UUID.randomUUID().toString().substring(0, 10));

        JSONObject parameter = new JSONObject(); // parameter参数
        JSONObject chat = new JSONObject();
        chat.set("domain", "generalv3.5");
        chat.set("temperature", 0.5);
        chat.set("max_tokens", 4096);
        parameter.set("chat", chat);

        JSONObject payload = new JSONObject(); // payload参数
        JSONObject message = new JSONObject();
        JSONArray text = new JSONArray();

        for (GptRespDto.RoleContent tempRoleContent : historyList) {
            text.add(JSONUtil.toJsonStr(tempRoleContent));
        }

        // 最新问题
        GptRespDto.RoleContent roleContent = new GptRespDto.RoleContent();
        roleContent.setRole("user");
        roleContent.setContent(question);
        text.add(JSONUtil.toJsonStr(roleContent));
        historyList.add(roleContent);

        message.set("text", text);
        payload.set("message", message);

        requestJson.set("header", header);
        requestJson.set("parameter", parameter);
        requestJson.set("payload", payload);
        log.debug("requestJson: {}", requestJson);
        boolean send = webSocket.send(requestJson.toString());
        log.debug("发送消息结果: {}", send);
    }

    private void waitAnswer(SparkWebSocketMessageListener sparkWebSocketMessageListener) {
        Lock lock = sparkWebSocketMessageListener.getLock();
        Condition condition = sparkWebSocketMessageListener.getCondition();

        lock.lock();
        try {
            condition.await();
        } catch (InterruptedException e) {
            throw new RuntimeException(e);
        } finally {
            lock.unlock();
        }
        // wait/notify方式
//        synchronized (totalFlag) {
//            try {
//                totalFlag.wait();
//            } catch (InterruptedException e) {
//                throw new RuntimeException(e);
//            }
//        }
    }
}
